#include "../server/logging/Logging.h"
#include "../server/base/Utils.h"
#include "../server/net/EventLoop.h"
#include "../server/net/EventLoopThreadPool.h"
#include "../server/net/TcpClient.h"
#include "../server/net/TcpServer.h"
#include "../server/thread/ThreadLocalSingleton.h"
#include "./protorpc/RpcCodec.h"
#include "./protobuf/ProtobufCodecLite.h"
#include "./proto/rpc.pb.h"

#include <stdio.h>
#include <unistd.h>

using namespace yhaida;

struct RawMessage
{
    RawMessage(std::string m)
        : message(std::move(m)),
          _id(0),
          _loc(NULL)
    {
    }

    uint64_t id() const { return _id; }
    void set_id(uint64_t x) { _id = x; }

    bool parse(const std::string &tag);
    void updateId();

public:
    std::string message;

private:
    uint64_t _id;
    const void *_loc;
};

class BackendSession : Utils::nocapyable
{
public:
    BackendSession(EventLoop *loop, const InetAddress &serverAddr);
    void connect()
    {
        _client.connect();
    }
    template <typename MSG>
    // bool send(rpc::RpcMessage &msg, const TcpConnectionPtr &clientConn)
    bool send(MSG &msg, const TcpConnectionPtr &clientConn)
    {
        _loop->assertInLoopThread();
        if (_conn)
        {
            //为了保证序号的唯一性，将原来的msg.id()先保存后然后再附上一个唯一的id
            uint64_t id = ++_nextId;
            Request r = {msg.id(), clientConn};
            assert(_outstandings.find(id) == _outstandings.end());
            _outstandings[id] = r;
            msg.set_id(id);
            // _codec.send(_conn, msg);
            sendTo(_conn, msg);
            return true;
        }
        else
        {
            return false;
        }
    }

private:
    void sendTo(const TcpConnectionPtr &conn, const rpc::RpcMessage &msg)
    {
        _codec.send(conn, msg);
    }

    void sendTo(const TcpConnectionPtr &conn, RawMessage &msg)
    {
        msg.updateId();
        conn->send(msg.message);
    }

    void onConnection(const TcpConnectionPtr &conn);

    void onRpcMessage(const TcpConnectionPtr &conn,
                      const rpc::RpcMessagePtr msg,
                      Timestamp);

    bool onRawMessage(const TcpConnectionPtr &,
                      std::string message,
                      Timestamp);

    template <typename MSG>
    void onMessageT(MSG &msg)
    {
        _loop->assertInLoopThread();
        std::map<uint64_t, Request>::iterator it = _outstandings.find(msg.id());
        if (it != _outstandings.end())
        {
            uint64_t origId = it->second.origId;
            TcpConnectionPtr clientConn = it->second.clientConn.lock();
            _outstandings.erase(it);

            if (clientConn)
            {
                // LOG_DEBUG << "send back " << origId << " of " << clientConn->name()
                //           << " using " << msg.id() << " from " << conn_->name();
                msg.set_id(origId);
                sendTo(clientConn, msg);
            }
        }
        else
        {
            // LOG_ERROR
        }
    }

private:
    struct Request
    {
        uint64_t origId;
        std::weak_ptr<TcpConnection> clientConn;
    };

private:
    EventLoop *_loop;
    TcpClient _client;
    rpc::RpcCodec _codec;
    TcpConnectionPtr _conn; //_client的TCP连接
    uint64_t _nextId;
    std::map<uint64_t, Request> _outstandings;
};

class Balancer : Utils::nocapyable
{
public:
    Balancer(EventLoop *loop,
             const InetAddress &listenAddr,
             const std::vector<InetAddress> &backends);
    ~Balancer()
    {
    }
    void setThreadNum(int numThreads)
    {
        _server.setThreadNum(numThreads);
    }

    void start()
    {
        _server.start();
    }

private:
    struct PerThread
    {
        size_t current;
        std::vector<std::unique_ptr<BackendSession>> backends;
        PerThread() : current(0) {}
    };

    void initPerThread(EventLoop *ioLoop);

    void onConnection(const TcpConnectionPtr &conn);

    void onRpcMessage(const TcpConnectionPtr &conn,
                      const rpc::RpcMessagePtr &msg,
                      Timestamp);

    bool onRawMessage(const TcpConnectionPtr &conn,
                      std::string message,
                      Timestamp);
    template <typename MSG>
    void onMessageT(const TcpConnectionPtr &conn, MSG &msg)
    {
        bool succeed = false;
        for (size_t i = 0; i < _t_backends::instance().backends.size() && !succeed; ++i)
        {
            succeed = _t_backends::instance().backends[_t_backends::instance().current]->send(msg, conn);
            _t_backends::instance().current = (_t_backends::instance().current + 1) % _t_backends::instance().backends.size();
        }
        if (!succeed)
        {
            // TODO
        }
    }

private:
    typedef ThreadLocalSingleton<PerThread> _t_backends;
    TcpServer _server;
    rpc::RpcCodec _codec;
    std::vector<InetAddress> _backends;
    AtomicInt32 _threadCount;
};
